import os

import mmcv
# using a pre-trained detector. 预训练配置
from mmcv import Config


# root_path=r"/project/train/src_repo/mmdetection/" # 要改的

root_path=r"/home/deepin/Documents/ji_pingtai/mmdetection/" # 要改的

cfg_path=root_path+'configs/yolo/yolov3_d53_fp16_mstrain-608_273e_coco.py'

cfg = Config.fromfile(cfg_path)

from mmdet.apis import set_random_seed, inference_detector, show_result_pyplot

# Modify dataset type and path
cfg.dataset_type = 'CocoDataset' # 自定义的数据集的 类
cfg.data_root = root_path+'data/coco/'

# 训练集
cfg.data.train.type = 'CocoDataset'
cfg.data.train.data_root = cfg.data_root 
cfg.data.train.ann_file = 'annotations/instances_train2017.json' # 存 图片名字
cfg.data.train.img_prefix = 'train2017'# 图片位置

#验证集
cfg.data.val.type = 'CocoDataset'
cfg.data.val.data_root = cfg.data_root 
cfg.data.val.ann_file = 'annotations/instances_val2017.json'
cfg.data.val.img_prefix = 'val2017'

# 测试集
cfg.data.test.type = 'CocoDataset'
cfg.data.test.data_root =cfg.data_root 
cfg.data.val.ann_file = 'annotations/instances_val2017.json'
cfg.data.val.img_prefix = 'val2017'



cfg.device='cuda' # 'ConfigDict' object has no attribute 'device'

# modify num classes of the model in box head 类别数
cfg.model.bbox_head.num_classes = 2 #----类别数

# If we need to finetune a model based on a pre-trained detector, we need to
# use load_from to set the path of checkpoints. 预训练模型 
#下载： wget http://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth

cfg.load_from = root_path+'checkpoints/yolov3_d53_fp16_mstrain-608_273e_coco_20210517_213542-4bc34944.pth'


# Set up working dir to save files and logs. 保存的文件夹
cfg.work_dir = 'work_dir'

# The original learning rate (LR) is set for 8-GPU training.
# We divide it by 8 since we only use one GPU.
#原始学习率（LR）被设置用于8-GPU训练。
#我们将其除以8，因为我们只使用一个GPU。
# cfg.optimizer.lr = 0.05 / 8 # 过大，学习率过高，梯度下降快，梯度消失
# cfg.lr_config.warmup = None
cfg.log_config.interval = 20 #打印log 间隔

# Change the evaluation metric since we use customized dataset.
#更改评估指标，因为我们使用自定义数据集。
# cfg.evaluation.metric = 'mAP'

# We can set the evaluation interval to reduce the evaluation times

cfg.evaluation.interval = 10 # 评估间隔，以减少评估时间

cfg.evaluation.save_best='auto' # 保存最佳best的pth



# We can set the checkpoint saving interval to reduce the storage cost

cfg.checkpoint_config.interval = 25  # 保存间隔:  设置检查点：，以降低存储成本

# cfg.runner = dict(type='EpochBasedRunner', max_epochs=20)# ---训练 轮次


# Set seed thus the results are more reproducible #播种，结果更具可重复性
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

# We can also use tensorboard to log the training process #我们还可以使用tensorboard记录培训过程
cfg.log_config.hooks = [
    dict(type='TextLoggerHook'),
    dict(type='TensorboardLoggerHook')]


# We can initialize the logger for training and have a look
# at the final config used for training
#我们可以初始化记录器进行培训并查看
#在用于培训的最终配置中
print(f'Config:\n{cfg.pretty_text}')



#训练----------------------------------------------------
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector


# Build dataset 构建 数据集
print(cfg.data.train)
datasets = [build_dataset(cfg.data.train)]

# Build the detector 构建 识别
model = build_detector(cfg.model)

# Add an attribute for visualization convenience 添加属性以方便可视化，模型的类别
model.CLASSES = datasets[0].CLASSES

# Create work_dir 文件夹
mmcv.mkdir_or_exist(os.path.abspath(cfg.work_dir))

#保存 cfg 配置：
# 把自定义的 config 在 自定义的working dir 下 保存一份------------推理和测试 可以用到的
cfg.dump(F'{cfg.work_dir}/yolox_cfgformat.py')

train_detector(model, datasets, cfg, distributed=False, validate=True) # -------进行----训练-----


# # # 查看 评估结果 load tensorboard in colab
# # %load_ext tensorboard
# #
# # # see curves in tensorboard
# # tensorboard --logdir ./tutorial_exps  #--------运行这个看看 数的 曲线



